{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Qwen2大模型微调入门-命名实体识别任务\n",
    "\n",
    "Qwen2是通义千问团队最近开源的大语言模型，由阿里云通义实验室研发。\n",
    "\n",
    "以Qwen2作为基座大模型，通过指令微调的方式做高精度的命名实体识别（NER），是学习入门LLM微调、建立大模型认知的非常好的任务。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在本文中，我们会使用 Qwen2-1.5b-Instruct 模型在 中文NER 数据集上做指令微调训练，同时监控训练过程、评估模型效果。\n",
    "\n",
    "* 显存要求：10GB左右"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.环境安装\n",
    "\n",
    "本案例基于Python>=3.8，请在您的计算机上安装好Python，并且有一张英伟达显卡（显存要求并不高，大概10GB左右就可以跑）。\n",
    "\n",
    "我们需要安装以下这几个Python库，在这之前，请确保你的环境内已安装好了pytorch以及CUDA：\n",
    "本案例测试于modelscope==1.14.0、transformers>=4.41.2,<=4.45.0、datasets>=2.16.0,<=2.21.0、peft>=0.11.1,<=0.12.0、accelerate>=0.30.1,<=0.34.2、pandas、matplotlib、tensorboard、tensorboardX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install modelscope==1.14.0 transformers>=4.41.2,<=4.45.0 datasets>=2.16.0,<=2.21.0 peft>=0.11.1,<=0.12.0 pandas accelerate>=0.30.1,<=0.34.2 matplotlib tensorboard tensorboardX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 知识点1：什么是指令微调？\n",
    "大模型指令微调（Instruction Tuning）是一种针对大型预训练语言模型的微调技术，其核心目的是增强模型理解和执行特定指令的能力，使模型能够根据用户提供的自然语言指令准确、恰当地生成相应的输出或执行相关任务。\n",
    "\n",
    "指令微调特别关注于提升模型在遵循指令方面的一致性和准确性，从而拓宽模型在各种应用场景中的泛化能力和实用性。\n",
    "\n",
    "在实际应用中，我的理解是，指令微调更多把LLM看作一个更智能、更强大的传统NLP模型（比如Bert），来实现更高精度的NLP任务。所以这类任务的应用场景覆盖了以往NLP模型的场景，甚至很多团队拿它来标注互联网数据。\n",
    "\n",
    "知识点2：什么是命名实体识别？\n",
    "命名实体识别 (NER) 是一种NLP技术，主要用于识别和分类文本中提到的重要信息（关键词）。这些实体可以是人名、地名、机构名、日期、时间、货币值等等。 NER 的目标是将文本中的非结构化信息转换为结构化信息，以便计算机能够更容易地理解和处理。\n",
    "\n",
    "![image.png](images/ner.png)\n",
    "\n",
    "NER 也是一项非常实用的技术，包括在互联网数据标注、搜索引擎、推荐系统、知识图谱、医疗保健等诸多领域有广泛应用。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.准备数据集\n",
    "本案例使用的是HuggingFace上的chinese_ner_sft数据集，该数据集主要被用于训练命名实体识别模型。\n",
    "\n",
    "![image.png](images/dataset.png)\n",
    "\n",
    "chinese_ner_sft由不同来源、不同类型的几十万条数据组成，应该是我见过收录最齐全的中文NER数据集。\n",
    "\n",
    "这次训练我们不需要用到它的全部数据，只取其中的CCFBDCI数据集（中文命名实体识别算法鲁棒性评测数据集）进行训练，该数据集包含LOC（地点）、GPE（地理）、ORG（组织）和PER（人名）四种实体类型标注，每条数据的例子如下：\n",
    "\n",
    "```\n",
    "{\n",
    "  \"text\": \"今天亚太经合组织第十二届部长级会议在这里开幕，中国外交部部长唐家璇、外经贸部部长石广生出席了会议。\",\n",
    "  \"entities\": [\n",
    "    {\n",
    "        \"start_idx\": 23,\n",
    "        \"end_idx\": 25,\n",
    "        \"entity_text\": \"中国\",\n",
    "        \"entity_label\": \"GPE\",\n",
    "        \"entity_names\": [\"地缘政治实体\", \"政治实体\", \"地理实体\", \"社会实体\"]},\n",
    "        {\n",
    "            \"start_idx\": 25,\n",
    "            \"end_idx\": 28,\n",
    "            \"entity_text\": \"外交部\",\n",
    "            \"entity_label\": \"ORG\",\n",
    "            \"entity_names\": [\"组织\", \"团体\", \"机构\"]\n",
    "        },\n",
    "        {\n",
    "            \"start_idx\": 30,\n",
    "            \"end_idx\": 33,\n",
    "            \"entity_text\": \"唐家璇\",\n",
    "            \"entity_label\": \"PER\",\n",
    "            \"entity_names\": [\"人名\", \"姓名\"]\n",
    "        }, \n",
    "        ...\n",
    "    ],\n",
    "\"data_source\": \"CCFBDCI\"\n",
    "}\n",
    "```\n",
    "\n",
    "其中text是输入的文本，entities是文本抽取出的实体。我们的目标是希望微调后的大模型能够根据由text组成的提示词，预测出一个json格式的实体信息：\n",
    "\n",
    "```\n",
    "输入：今天亚太经合组织第十二届部长级会议在这里开幕，中国外交部部长唐家璇、外经贸部部长石广生出席了会议。\n",
    "\n",
    "大模型输出：{\"entity_text\":\"中国\", \"entity_label\":\"组织\"}{\"entity_text\":\"唐家璇\", \"entity_label\":\"人名\"}...\n",
    "```\n",
    "\n",
    "##### 2.1 现在我们将数据集下载到本地目录。下载方式是前往chinese_ner_sft - huggingface下载ccfbdci.jsonl到项目根目录下即可：\n",
    "\n",
    "![image.png](images/dataset_download.png)\n",
    "\n",
    "##### 2.2 将ccfbdci.jsonl进行处理，转换成new_train.jsonl和new_test.jsonl"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2.将train.jsonl和test.jsonl进行处理，转换成new_train.jsonl和new_test.jsonl\n",
    "\n",
    "import json\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "def dataset_jsonl_transfer(origin_path, new_path):\n",
    "    \"\"\"\n",
    "    将原始数据集转换为大模型微调所需数据格式的新数据集\n",
    "    \"\"\"\n",
    "    messages = []\n",
    "\n",
    "    # 读取旧的JSONL文件\n",
    "    with open(origin_path, \"r\") as file:\n",
    "        for line in file:\n",
    "            # 解析每一行的json数据\n",
    "            data = json.loads(line)\n",
    "            input_text = data[\"text\"]\n",
    "            entities = data[\"entities\"]\n",
    "            match_names = [\"地点\", \"人名\", \"地理实体\", \"组织\"]\n",
    "            \n",
    "            entity_sentence = \"\"\n",
    "            for entity in entities:\n",
    "                entity_json = dict(entity)\n",
    "                entity_text = entity_json[\"entity_text\"]\n",
    "                entity_names = entity_json[\"entity_names\"]\n",
    "                for name in entity_names:\n",
    "                    if name in match_names:\n",
    "                        entity_label = name\n",
    "                        break\n",
    "                \n",
    "                entity_sentence += f\"\"\"{{\"entity_text\": \"{entity_text}\", \"entity_label\": \"{entity_label}\"}}\"\"\"\n",
    "            \n",
    "            if entity_sentence == \"\":\n",
    "                entity_sentence = \"没有找到任何实体\"\n",
    "            \n",
    "            message = {\n",
    "                \"instruction\": \"\"\"你是一个文本实体识别领域的专家，你需要从给定的句子中提取 地点; 人名; 地理实体; 组织 实体. 以 json 格式输出, 如 {\"entity_text\": \"南京\", \"entity_label\": \"地理实体\"} 注意: 1. 输出的每一行都必须是正确的 json 字符串. 2. 找不到任何实体时, 输出\"没有找到任何实体\". \"\"\",\n",
    "                \"input\": f\"文本:{input_text}\",\n",
    "                \"output\": entity_sentence,\n",
    "            }\n",
    "            \n",
    "            messages.append(message)\n",
    "\n",
    "    # 保存重构后的JSONL文件\n",
    "    with open(new_path, \"w\", encoding=\"utf-8\") as file:\n",
    "        for message in messages:\n",
    "            file.write(json.dumps(message, ensure_ascii=False) + \"\\n\")\n",
    "\n",
    "\n",
    "# 加载、处理数据集和测试集\n",
    "train_dataset_path = \"ccfbdci.jsonl\"\n",
    "train_jsonl_new_path = \"ccf_train.jsonl\"\n",
    "\n",
    "if not os.path.exists(train_jsonl_new_path):\n",
    "    dataset_jsonl_transfer(train_dataset_path, train_jsonl_new_path)\n",
    "\n",
    "total_df = pd.read_json(train_jsonl_new_path, lines=True)\n",
    "train_df = total_df[int(len(total_df) * 0.1):]  # 取90%的数据做训练集\n",
    "test_df = total_df[:int(len(total_df) * 0.1)].sample(n=20)  # 随机取10%的数据中的20条做测试集"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. 下载/加载模型和tokenizer\n",
    "\n",
    "这里我们使用modelscope下载Qwen2-1.5B-Instruct模型（modelscope在国内，所以直接用下面的代码自动下载即可，不用担心速度和稳定性问题），然后把它加载到Transformers中进行训练："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from modelscope import snapshot_download, AutoTokenizer\n",
    "from transformers import AutoModelForCausalLM\n",
    "import os\n",
    "import torch\n",
    "\n",
    "\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # 设置使用的GPU编号\n",
    "\n",
    "model_id = \"qwen/Qwen2-1.5B-Instruct\"    \n",
    "model_dir = \"./qwen/Qwen2-1___5B-Instruct\"\n",
    "\n",
    "# 在modelscope上下载Qwen模型到本地目录下\n",
    "model_dir = snapshot_download(model_id, cache_dir=\"./\", revision=\"master\")\n",
    "\n",
    "# Transformers加载模型权重\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_dir, use_fast=False, trust_remote_code=True)\n",
    "model = AutoModelForCausalLM.from_pretrained(model_dir, device_map=\"auto\", torch_dtype=torch.bfloat16, trust_remote_code=True, use_cache=False)\n",
    "model.enable_input_require_grads()  # 开启梯度检查点时，要执行该方法"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. 预处理训练数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_func(example):\n",
    "    \"\"\"\n",
    "    将数据集进行预处理, 处理成模型可以接受的格式\n",
    "    \"\"\"\n",
    "\n",
    "    MAX_LENGTH = 384 \n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "    system_prompt = \"\"\"你是一个文本实体识别领域的专家，你需要从给定的句子中提取 地点; 人名; 地理实体; 组织 实体. 以 json 格式输出, 如 {\"entity_text\": \"南京\", \"entity_label\": \"地理实体\"} 注意: 1. 输出的每一行都必须是正确的 json 字符串. 2. 找不到任何实体时, 输出\"没有找到任何实体\".\"\"\"\n",
    "    \n",
    "    instruction = tokenizer(\n",
    "        f\"<|im_start|>system\\n{system_prompt}<|im_end|>\\n<|im_start|>user\\n{example['input']}<|im_end|>\\n<|im_start|>assistant\\n\",\n",
    "        add_special_tokens=False,\n",
    "    )\n",
    "    response = tokenizer(f\"{example['output']}\", add_special_tokens=False)\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    attention_mask = (\n",
    "        instruction[\"attention_mask\"] + response[\"attention_mask\"] + [1]\n",
    "    )\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    if len(input_ids) > MAX_LENGTH:  # 做一个截断\n",
    "        input_ids = input_ids[:MAX_LENGTH]\n",
    "        attention_mask = attention_mask[:MAX_LENGTH]\n",
    "        labels = labels[:MAX_LENGTH]\n",
    "    return {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"labels\": labels}   "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import Dataset\n",
    "\n",
    "train_ds = Dataset.from_pandas(train_df)\n",
    "train_dataset = train_ds.map(process_func, remove_columns=train_ds.column_names)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. 设置LORA"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from peft import LoraConfig, TaskType, get_peft_model\n",
    "\n",
    "config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM,\n",
    "    target_modules=[\n",
    "        \"q_proj\",\n",
    "        \"k_proj\",\n",
    "        \"v_proj\",\n",
    "        \"o_proj\",\n",
    "        \"gate_proj\",\n",
    "        \"up_proj\",\n",
    "        \"down_proj\",\n",
    "    ],\n",
    "    inference_mode=False,  # 训练模式\n",
    "    r=8,  # Lora 秩\n",
    "    lora_alpha=32,  # Lora alaph，具体作用参见 Lora 原理\n",
    "    lora_dropout=0.1,  # Dropout 比例\n",
    ")\n",
    "\n",
    "model = get_peft_model(model, config)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 6. 微调训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from transformers import TrainingArguments, Trainer, DataCollatorForSeq2Seq\n",
    "\n",
    "# Set up logging to a file\n",
    "logging.basicConfig(filename='training.log', level=logging.INFO)\n",
    "\n",
    "args = TrainingArguments(\n",
    "    output_dir=\"./output/Qwen2-NER\",\n",
    "    per_device_train_batch_size=4,\n",
    "    per_device_eval_batch_size=4,\n",
    "    gradient_accumulation_steps=4,\n",
    "    logging_steps=10,\n",
    "    num_train_epochs=2,\n",
    "    save_steps=100,\n",
    "    learning_rate=1e-4,\n",
    "    save_on_each_node=True,\n",
    "    gradient_checkpointing=True,\n",
    "    report_to=[\"tensorboard\"],  # Log to file\n",
    "    logging_dir='./logs',  # Directory for logging\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=train_dataset,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    ")\n",
    "\n",
    "trainer.train()\n",
    "\n",
    "\n",
    "# ====== 训练结束后的预测 ===== #\n",
    "\n",
    "def predict(messages, model, tokenizer):\n",
    "    device = \"cuda\"\n",
    "    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "    model_inputs = tokenizer([text], return_tensors=\"pt\").to(device)\n",
    "    generated_ids = model.generate(model_inputs.input_ids, max_new_tokens=512)\n",
    "    generated_ids = [\n",
    "        output_ids[len(input_ids) :]\n",
    "        for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
    "    ]\n",
    "\n",
    "    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
    "    print(response)\n",
    "\n",
    "    return response\n",
    "    \n",
    "\n",
    "test_text_list = []\n",
    "for index, row in test_df.iterrows():\n",
    "    instruction = row[\"instruction\"]\n",
    "    input_value = row[\"input\"]\n",
    "\n",
    "    messages = [\n",
    "        {\"role\": \"system\", \"content\": f\"{instruction}\"},\n",
    "        {\"role\": \"user\", \"content\": f\"{input_value}\"},\n",
    "    ]\n",
    "\n",
    "    response = predict(messages, model, tokenizer)\n",
    "    messages.append({\"role\": \"assistant\", \"content\": f\"{response}\"})\n",
    "    result_text = f\"{messages[0]}\\n\\n{messages[1]}\\n\\n{messages[2]}\"\n",
    "    test_text_list.append(result_text)\n",
    "\n",
    "# Log the predictions to a file\n",
    "with open('./logs/predictions.log', 'w', encoding='utf-8') as log_file:\n",
    "    for text in test_text_list:\n",
    "        log_file.write(text + \"\\n\\n\")\n",
    "\n",
    "print(\"Predictions have been logged to predictions.log\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext tensorboard\n",
    "\n",
    "%tensorboard --logdir logs"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
